import numpy as np

import CalcRule as CR
from object import object, orient_object
from PPO import ActorCritic
import torch


class field():
    def __init__(self, arg):
        self.arg = arg
        self.numA = arg.num1
        self.numB = arg.num2
        self.length = arg.field_length
        self.width = arg.field_width
        self.radius_player = arg.radius_player
        self.radius_soccer = arg.radius_soccer
        self.time_step = arg.time_step
        self.gamma = arg.gamma_velocity
        self.orient = arg.orient
        self.gate_length = arg.gate_length

        self.teamA = []
        self.teamB = []
        self.soccer = None
        # self.derive_state = self.derive_arc if arg.state_term == 1 else self.derive_pos
        # self.derive_state = self.derive_abs_pos
        self.derive_state = self.derive_pos
        # self.derive_state = self.derive_arc

    def derive_abs_state(self):
        state = []
        flag = isinstance(self.teamA[0].coord, object)
        if not self.orient:
            for i in range(self.numA):
                state += [self.teamA[i].coord.x, self.teamA[i].coord.y, self.teamA[i].vel.x, self.teamA[i].vel.y]
            for i in range(self.numB):
                state += [self.teamB[i].coord.x, self.teamB[i].coord.y, self.teamB[i].vel.x, self.teamB[i].vel.y]
            state += [self.soccer.coord.x, self.soccer.coord.y, self.soccer.vel.x, self.soccer.vel.y]
        else:
            for i in range(self.numA):
                state += [self.teamA[i].coord.x, self.teamA[i].coord.y, self.teamA[i].coord.phi, self.teamA[i].vel.x,
                          self.teamA[i].vel.y, self.teamA[i].vel.phi]
            for i in range(self.numB):
                state += [self.teamB[i].coord.x, self.teamB[i].coord.y, self.teamB[i].coord.phi, self.teamB[i].vel.x,
                          self.teamB[i].vel.y, self.teamB[i].vel.phi]
            state += [self.soccer.coord.x, self.soccer.coord.y, self.soccer.coord.phi, self.soccer.vel.x,
                      self.soccer.vel.y, self.soccer.vel.phi]
        return np.array(state)

    def derive_abs_pos(self):
        state = []
        flag = isinstance(self.teamA[0].coord, object)
        if not self.orient:
            for i in range(self.numA):
                state += [self.teamA[i].coord.x, self.teamA[i].coord.y]
            for i in range(self.numB):
                state += [self.teamB[i].coord.x, self.teamB[i].coord.y]
            state += [self.soccer.coord.x, self.soccer.coord.y]
            state = state + [self.soccer.vel.x, self.soccer.vel.y]
        else:
            for i in range(self.numA):
                state += [self.teamA[i].coord.x, self.teamA[i].coord.y, self.teamA[i].coord.phi]
            for i in range(self.numB):
                state += [self.teamB[i].coord.x, self.teamB[i].coord.y, self.teamB[i].coord.phi]
            state += [self.soccer.coord.x, self.soccer.coord.y, self.soccer.coord.phi]
        return np.array(state)

    def derive_pos(self):
        state = []
        for i in range(self.numA):
            state.append(-self.teamA[i].coord.x + self.soccer.coord.x)
            state.append(-self.teamA[i].coord.y + self.soccer.coord.y)
        for i in range(self.numB):
            state.append(-self.teamB[i].coord.x + self.soccer.coord.x)
            state.append(-self.teamB[i].coord.y + self.soccer.coord.y)
        state.append(self.width / 2 - self.soccer.coord.x)
        state.append(-self.soccer.coord.y)
        state += [self.soccer.vel.x, self.soccer.vel.y]
        return np.array(state)

    def derive_arc(self):
        state = []
        for i in range(self.numA):
            dx = -self.teamA[i].coord.x + self.soccer.coord.x
            dy = self.teamA[i].coord.y - self.soccer.coord.y
            r, theta = CR.arc(dx, dy)
            state = state + [4 * r / self.width, theta]
        for i in range(self.numB):
            dx = self.teamB[i].coord.x - self.soccer.coord.x
            dy = self.teamB[i].coord.y - self.soccer.coord.y
            r, theta = CR.arc(dx, dy)
            state = state + [4 * r / self.width, theta]
        dx = self.width / 2 - self.soccer.coord.x
        dy = -self.soccer.coord.y
        r, theta = CR.arc(dx, dy)
        state = state + [4 * r / self.width, theta]
        state = state + [self.soccer.vel.x, self.soccer.vel.y]
        return np.array(state)

    def reset(self, area, vel_area):
        xlimA, ylimA, xlimB, ylimB, xball, yball = area
        vxlimA, vylimA, vxlimB, vylimB, vxball, vyball = vel_area
        while True:
            self.teamA = []
            self.teamB = []
            if not self.orient:
                for i in range(self.numA):
                    random_coord = CR.random(xlimA, ylimA)
                    random_vel = CR.random(vxlimA, vylimA)
                    self.teamA.append(object(random_coord, random_vel, self.radius_player, i))
                for i in range(self.numB):
                    random_coord = CR.random(xlimB, ylimB)
                    random_vel = CR.random(vxlimB, vylimB)
                    self.teamB.append(object(random_coord, random_vel, self.radius_player, i))
                random_coord = CR.random(xball, yball)
                random_vel = CR.random(vxball, vyball)
                self.soccer = object(random_coord, random_vel, self.radius_soccer)
            else:
                for i in range(self.numA):
                    random_coord = CR.orient_random(xlimA, ylimA)
                    random_vel = CR.orient_random(vxlimA, vylimA)
                    self.teamA.append(orient_object(random_coord, random_vel, self.radius_player, i))
                for i in range(self.numB):
                    random_coord = CR.orient_random(xlimB, ylimB)
                    random_vel = CR.orient_random(vxlimB, vylimB)
                    self.teamB.append(orient_object(random_coord, random_vel, self.radius_player, i))
                random_coord = CR.orient_random(xball, yball)
                random_vel = CR.orient_random(vxball, vyball)
                self.soccer = orient_object(random_coord, random_vel, self.radius_soccer)
            if not self.collide():
                state = self.derive_state()
                break
        return state

    def collide(self):
        for i in range(self.numA):
            for j in range(i + 1, self.numA):

                dist_to_other = self.teamA[i].coord.dist(self.teamA[j].coord)
                if dist_to_other <= 2 * self.radius_player:
                    return True
            for j in range(0, self.numB):
                dist_to_other = self.teamA[i].coord.dist(self.teamB[j].coord)
                if dist_to_other <= 2 * self.radius_player:
                    return True
        for i in range(self.numB):
            for j in range(i + 1, self.numB):
                dist_to_other = self.teamB[i].coord.dist(self.teamB[j].coord)
                if dist_to_other <= 2 * self.radius_player:
                    return True
        return False

    def detect_player(self):
        for i in range(self.numA):
            # boundary
            if abs(self.teamA[i].coord.x) > self.width / 2 - self.radius_player:
                self.teamA[i].coord.x = (self.width / 2 - self.radius_player) * self.teamA[i].coord.x / abs(
                    self.teamA[i].coord.x)
            if abs(self.teamA[i].coord.y) > self.length / 2 - self.radius_player:
                self.teamA[i].coord.y = (self.length / 2 - self.radius_player) * self.teamA[i].coord.y / abs(
                    self.teamA[i].coord.y)

            # collide with each other
            for j in range(i + 1, self.numA):
                dist_to_other = self.teamA[i].coord.dist(self.teamA[j].coord)
                if dist_to_other <= 1.5 * self.radius_player:
                    return True
                if dist_to_other <= 2 * self.radius_player:
                    self.teamA[i].crush(self.teamA[j])
            for j in range(0, self.numB):
                dist_to_other = self.teamA[i].coord.dist(self.teamB[j].coord)
                if dist_to_other <= 1.5 * self.radius_player:
                    return True
                if dist_to_other <= 2 * self.radius_player:
                    self.teamA[i].crush(self.teamB[j])

            # collide with ball
            dist_to_ball = self.teamA[i].coord.dist(self.soccer.coord)
            if dist_to_ball < self.radius_player:
                return True
            if dist_to_ball <= self.radius_player + self.radius_soccer:
                self.soccer.strike(self.teamA[i])
        for i in range(self.numB):
            if abs(self.teamB[i].coord.x) > self.width / 2 - self.radius_player:
                self.teamB[i].coord.x = (self.width / 2 - self.radius_player) * self.teamB[i].coord.x / abs(
                    self.teamB[i].coord.x)
            if abs(self.teamB[i].coord.y) > self.length / 2 - self.radius_player:
                self.teamB[i].coord.y = (self.length / 2 - self.radius_player) * self.teamB[i].coord.y / abs(
                    self.teamB[i].coord.y)
            for j in range(i + 1, self.numB):
                dist_to_other = self.teamB[i].coord.dist(self.teamB[j].coord)
                if dist_to_other <= 1.5 * self.radius_player:
                    return True
                if dist_to_other <= 2 * self.radius_player:
                    self.teamB[i].crush(self.teamB[j])
            dist_to_ball = self.teamB[i].coord.dist(self.soccer.coord)
            if dist_to_ball < self.radius_player:
                return True
            if dist_to_ball <= self.radius_player + self.radius_soccer:
                self.soccer.strike(self.teamB[i])
        return False

    def set_vel(self, command):
        if not self.orient:
            for i in range(self.numA):
                self.teamA[i].vel = CR.vec2D(command[2 * i], command[2 * i + 1])
            for i in range(self.numB):
                self.teamB[i].vel = CR.vec2D(command[2 * self.numA + 2 * i], command[2 * self.numA + 2 * i + 1])
        else:
            for i in range(self.numA):
                self.teamA[i].vel = CR.orient_vec2D(command[3 * i], command[3 * i + 1], command[3 * i + 2])
            for i in range(self.numB):
                self.teamB[i].vel = CR.orient_vec2D(command[3 * self.numA + 3 * i], command[3 * self.numA + 3 * i + 1],
                                                    command[3 * self.numA + 3 * i + 2])

    def detect_soccer(self):
        if (self.soccer.coord.y > self.length / 2 - self.soccer.radius and self.soccer.vel.y > 0) or (
                self.soccer.coord.y < -self.length / 2 + self.soccer.radius and self.soccer.vel.y < 0):
            self.soccer.soccer_bounce(False)
        if self.soccer.coord.x < -self.width / 2 + self.soccer.radius:
            if self.gate_length / 2 > self.soccer.coord.y > - self.gate_length / 2:
                return -1
            # else:
            if self.soccer.vel.x < 0:
                self.soccer.soccer_bounce(True)
        if self.soccer.coord.x > self.width / 2 - self.soccer.radius:
            if self.gate_length / 2 > self.soccer.coord.y > - self.gate_length / 2:
                return 1
            else:
                if self.soccer.vel.x > 0:
                    self.soccer.soccer_bounce(True)
        return 0

    def set_coord(self):
        for i in range(self.numA):
            self.teamA[i].process(False, self.time_step, self.gamma)
        for i in range(self.numB):
            self.teamB[i].process(False, self.time_step, self.gamma)
        flag = self.detect_soccer()
        self.soccer.process(True, self.time_step, self.gamma)
        return flag

    def vel_polar2eu(self, command):
        state = self.derive_arc()
        eu_command = command.copy()
        if self.orient:
            coef = 3
        else:
            coef = 2
        for i in range(self.numA):
            c = np.cos(state[coef * i + 1])
            s = np.sin(state[coef * i + 1])
            eu_command[coef * i] = command[coef * i + 1] * s - command[coef * i] * c
            eu_command[coef * i + 1] = command[coef * i + 1] * c + command[coef * i] * s
        for i in range(self.numA, self.numA + self.numB):
            c = np.cos(state[coef * i + 1])
            s = np.sin(state[coef * i + 1])
            eu_command[coef * i] = -command[coef * i + 1] * s + command[coef * i] * c
            eu_command[coef * i + 1] = command[coef * i + 1] * c + command[coef * i] * s
        return eu_command

    def defence_policy(self):
        return -np.array([0.1, 0] * self.numB)

    def step(self, action1, action2):
        state = self.derive_state()
        command = np.concatenate((action1, action2))
        command = self.vel_polar2eu(command)
        self.set_vel(command)
        r = 0
        bug = self.detect_player()
        flag = self.set_coord()
        state_ = self.derive_state()
        if bug:
            r = -1
            return state_, r, bug

        if flag == 0:
            r += np.linalg.norm(state[0:2]) - 0.01 * np.linalg.norm(state_[0:2])
            # r += 0.02 * np.linalg.norm(state[-4:-2]) - 0.02 * np.linalg.norm(state_[-4:-2])
            done = False
        else:
            # if flag == 1:
            #     r += 100
            # else:
            #     r -= 100
            done = True
        return state_, r, done

    def render(self, frames, save_dir=None):
        import visualize
        visualize.draw(frames, save_dir)
